import os
from pathlib import Path

import torch
from PIL import Image
from torchvision import datasets, transforms
from torch.utils.tensorboard import SummaryWriter

from Forecase import Forecase
from sort.data.DataHandler import DataHandler
from sort.nerve.NeuralNetwork import NeuralNetwork

if __name__ == '__main__':

    # 创建数据集
    data_dir = 'I:\\Workspace\\Dataset\\data_inception\\train'  # 你的数据集路径
    batch_size = 32

    # 数据预处理和增强
    transform = DataHandler(data_dir, batch_size).transform

    train_dataset = datasets.ImageFolder(data_dir, transform=transform)

    # 数据集拆分比例
    train_size = int(0.7 * len(train_dataset))
    val_size = len(train_dataset) - train_size

    # 随机拆分数据集
    train_dataset, val_dataset = torch.utils.data.random_split(train_dataset, [train_size, val_size])

    # 数据加载器
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size)

    for X, y in train_loader:
        print(f"Shape of X [N, C, H, W]: {X.shape}")
        print(f"Shape of y: {y.shape} {y.dtype}")
        break

    # 获取 cpu, gpu 或 mps 设备用于加速训练.
    device = (
        "cuda"
        if torch.cuda.is_available()  # 如果GPU可用，则使用GPU
        else "mps"
        if torch.backends.mps.is_available()  # 如果MPS可用，则使用MPS
        else "cpu"  # 否则使用CPU
    )
    # torch.backends.mps.is_available()

    print(f"Using {device} device")

    classes = train_dataset.dataset.classes

    # 加载模型结构和权重
    # model = torch.load('your_model.pth')

    # 或者加载模型权重到预先创建的模型
    model_path = 'model/ImageSort.pth'
    model = NeuralNetwork(len(classes))
    model.load_state_dict(torch.load(model_path))

    forecase = Forecase(model, transform)
    forecase.run_forecast('I:\\Workspace\\Dataset\\data_inception\\test\\5', classes)
